﻿using log4net;
using Newtonsoft.Json.Linq;
using per.esmool.utils;
using System;
using System.Collections.Generic;
using System.Linq;

namespace FlowSimulate {

    public class AIInvoker(string promptsFilePath) {

        private static readonly ILog LOG = LogManager.GetLogger(typeof(AIInvoker));

        public AIRequester AIRequester { get; set; }

        public string PromptFile { get; set; }
        public int MinAskTimes { get; set; }
        public int MaxAskTimes { get; set; }

        readonly Dictionary<string, List<string>> prompts = LoadPrompts(promptsFilePath);
        readonly RandKit rand = new();

        private static Dictionary<string, List<string>> LoadPrompts(string filepath) {
            var data = YamlKit.ParseFile<JArray>(filepath);
            data = JsonKit.Merge(data, "key", new() {
                MergeNullValueHandling = MergeNullValueHandling.Merge,
                MergeArrayHandling = MergeArrayHandling.Replace,
            });
            var qs = from item in data
                     where item is JObject
                     let json = item as JObject
                     let stage = JsonKit.GetValue<string>(json, "stage", "*")
                     let prompt = json.Value<string>("prompt")
                     select new { Stage = stage, Prompt = prompt };
            return qs
                .GroupBy(x => x.Stage, x => x.Prompt)
                .ToDictionary(x => x.Key, x => x.ToList());
        }

        private string GetAPrompt(string stageKey, HashSet<string> used) {
            var commonPrompts = this.prompts["*"];
            var stagePrompts = this.prompts.ContainsKey(stageKey) ? this.prompts[stageKey] : [];
            var prompts = commonPrompts.Concat(stagePrompts).Where(x => !used.Contains(x)).ToList();
            if (prompts.Count == 0) return null;
            return this.rand.Sample(prompts);
        }

        public string Invoke(string stageKey, string moduleName, string stageName, string taskName, string moduleDescription) {
            var used = new HashSet<string>();
            
            // HINT: 目前先不用多次重复, 不行了再换成多次重复, 多次重复的用途是构建机器人的上下文语境
            var desiredTrials = 1; // this.rand.SampleValue(this.MinAskTimes, this.MaxAskTimes);

            string result = null;
            for (int i=0; i<desiredTrials; i++) {
                var prompt = this.GetAPrompt(stageKey, used);
                var question = string.Format(prompt, moduleName, stageName, taskName, moduleDescription);
                result = this.RequestAI(question);
            }

            return result;
        }

        private string RequestAI(string question) {
            try {
                var accessToken = this.AIRequester.RequestAccessToken();
                var result = AIRequester.ERNIEBot(accessToken, question);
                return result;
            } catch (Exception err) {
                LOG.Warn($"Failed to request AI", err);
                return null;
            }
        }
    }
}